#!/usr/bin/python3
# coding=utf-8
import openai

from flask import request, jsonify
from flask import Flask
import json
from prompts import *
from flask_cors import cross_origin
client = openai.Client(
    api_key="cannot be empty",
    base_url=f"http://{xinference_url}:{xinference_port}/v1"
)

def chat_model(prompt, model_uid=model_uid, max_tokens=512, temperature=0.7):

    response = client.chat.completions.create(
        model=model_uid,
        messages=[
            {
                "content": prompt,
                "role": "user",
            }
        ],
        max_tokens=max_tokens,
        temperature=temperature
    )
    return response.choices[0].message.content

app = Flask(__name__)
app.debug = False

@app.route('/model/extract/topics', methods=['post'])
@cross_origin()
def post_http_topics():
    req = request.get_json()
    title = req.get('title', "")
    content = req.get("content", "")
    prompt = get_topics_prompt(title, content)
    response = chat_model(prompt)
    response = response.replace("领域", "").replace(":", "").replace("：", "")
    response = response.strip()
    try:
        data = json.loads(response)
    except:
        data = check_lis_string(topics, response)
    ret = {
        "code": 200,
        "data": data,
        "message": "success"
    }
    return jsonify(ret)
    # 返回JSON数据。

@app.route('/model/extract/domains', methods=['post'])
@cross_origin()
def post_http_domains():
    req = request.get_json()
    title = req.get('title', "")
    content = req.get("content", "")
    prompt = get_domains_prompt(title, content)
    response = chat_model(prompt)
    response = response.replace("主题", "").replace(":", "").replace("：", "")
    response = response.strip()
    try:
        data = json.loads(response)
    except:
        data = check_lis_string(domains, response)
        data = [d + "领域"for d in data]
    ret = {
        "code": 200,
        "data": data,
        "message": "success"
    }
    return jsonify(ret)



def check_lis_string(lis: list[str], string:str):
    for l in lis:
        if l in string:
            return [l]
    return ["其他"]




if __name__ == '__main__':
    app.run(host='0.0.0.0', port=6002)
    # 这里指定了地址和端口号。也可以不指定地址填0.0.0.0那么就会使用本机地址ip
